from functools import partial
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import Linear, BatchNorm1d, Sequential, ReLU
from torch_geometric.nn import global_mean_pool, global_add_pool, GINConv, GATConv
from gcn_conv import GCNConv
import random
import pdb


class CausalGCN(torch.nn.Module):
    """GCN with BN and residual connection."""
    # 定义了类的初始化方法，其中包括一系列参数，如输入特征的数量 num_features，输出类别的数量 num_classes，以及其他与模型结构相关的参数。
    def __init__(self, num_features,
                       num_classes, args,
                       gfn=False, 
                       collapse=False, 
                       residual=False,
                       res_branch="BNConvReLU", 
                       global_pool="sum", 
                       dropout=0, 
                       edge_norm=True):
        # super(CausalGCN, self).__init__(): 调用父类（torch.nn.Module）的初始化方法，确保正确地初始化模型。
        super(CausalGCN, self).__init__()
        # 从传入的参数 args 中获取图卷积层的数量。
        num_conv_layers = args.layers
        # 从传入的参数 args 中获取隐藏层的大小。
        hidden = args.hidden
        # 将传入的参数 args 存储为类的属性，以便在模型中的其他方法中使用。
        self.args = args
        # 将全局池化函数 global_add_pool 存储为类的属性，以便在模型的 forward 方法中使用。
        self.global_pool = global_add_pool
        # 将传入的 dropout 参数存储为类的属性
        self.dropout = dropout
        # 将 args 中的 with_random 参数存储为类的属性，用于控制是否包含随机特征。
        self.with_random = args.with_random
        # 将 args 中的 without_node_attention 参数存储为类的属性，用于控制是否包含节点注意力机制。
        self.without_node_attention = args.without_node_attention
        # 将 args 中的 without_edge_attention 参数存储为类的属性，用于控制是否包含边注意力机制。
        self.without_edge_attention = args.without_edge_attention
        # 使用 functools.partial 函数创建了一个偏函数 GConv，它基于 GCNConv 类，并固定了部分参数。
        # 其中 edge_norm 和 gfn 是传入的参数，用于控制是否进行边归一化和是否使用 Gated Fusion Network（GFN）
        GConv = partial(GCNConv, edge_norm=edge_norm, gfn=gfn)

        # 将输入特征的维度作为输入隐藏层的大小
        hidden_in = num_features
        # 将类别数量存储为类的属性
        self.num_classes = num_classes
        # 将类别数量作为输出隐藏层的大小
        hidden_out = num_classes
        # 将传入参数 args 中的全连接层数量存储为类的属性
        self.fc_num = args.fc_num
        # 创建一个一维批量归一化层 BatchNorm1d，并将其存储为类的属性。该层用于对输入特征进行批量归一化
        self.bn_feat = BatchNorm1d(hidden_in)
        # 创建了第一个图卷积层 GCNConv，用于对输入特征进行线性变换。它将输入特征的维度从 hidden_in 转换为 hidden。
        self.conv_feat = GCNConv(hidden_in, hidden, gfn=True) # linear transform
        # 创建了一个空的 PyTorch 模块列表 bns_conv，用于存储图卷积层后面的批量归一化层。
        self.bns_conv = torch.nn.ModuleList()
        # 创建了另一个空的 PyTorch 模块列表 convs，用于存储多个图卷积层。
        self.convs = torch.nn.ModuleList()
        # 遍历图卷积层
        for i in range(num_conv_layers):
            # 在 bns_conv 中添加一个批量归一化层，用于对当前图卷积层的输出进行归一化。
            self.bns_conv.append(BatchNorm1d(hidden))
            # 在 convs 中添加一个图卷积层，使用之前定义的偏函数 GConv 创建，其中隐藏层的输入和输出大小都是 hidden
            self.convs.append(GConv(hidden, hidden))

        # 这段代码主要是创建了用于计算注意力权重的线性层、批量归一化层以及上下文和对象节点的卷积层
        # 创建一个线性层，用于计算边注意力的权重，输入大小为 hidden * 2（因为连接的两个节点都经过了隐藏层），输出大小为 2（因为采用了 softmax 函数进行归一化，所以有两个输出）。
        self.edge_att_mlp = nn.Linear(hidden * 2, 2)
        # 创建另一个线性层，用于计算节点注意力的权重，输入大小为 hidden（因为每个节点都经过了隐藏层），输出大小为 2（同样采用 softmax 函数）。
        self.node_att_mlp = nn.Linear(hidden, 2)
        # 创建一个批量归一化层，用于对上下文节点的特征进行归一化。
        self.bnc = BatchNorm1d(hidden)
        # 创建另一个批量归一化层，用于对对象节点的特征进行归一化。
        self.bno= BatchNorm1d(hidden)
        # 创建上下文卷积层，用于处理上下文节点之间的信息传递。
        self.context_convs = GConv(hidden, hidden)
        # 创建对象卷积层，用于处理对象节点之间的信息传递。
        self.objects_convs = GConv(hidden, hidden)

        # 这些代码创建了用于上下文节点和对象节点的多层感知器（MLP），用于在信息传递之后对节点特征进行汇总。
        # 对于上下文节点，有两层线性层，分别是 fc1_c 和 fc2_c；对于对象节点，也有两层线性层，分别是 fc1_o 和 fc2_o。
        # 每一层后面都跟着一个批量归一化层
        # context mlp
        self.fc1_bn_c = BatchNorm1d(hidden)
        self.fc1_c = Linear(hidden, hidden)
        self.fc2_bn_c = BatchNorm1d(hidden)
        self.fc2_c = Linear(hidden, hidden_out)
        # object mlp
        self.fc1_bn_o = BatchNorm1d(hidden)
        self.fc1_o = Linear(hidden, hidden)
        self.fc2_bn_o = BatchNorm1d(hidden)
        self.fc2_o = Linear(hidden, hidden_out)
        # 这部分代码根据参数 args.cat_or_add 的值来决定使用拼接方式还是相加方式来处理上下文节点和对象节点的特征。根据不同的方式，创建不同的多层感知器结构。
        # random mlp
        if self.args.cat_or_add == "cat":
            self.fc1_bn_co = BatchNorm1d(hidden * 2)
            self.fc1_co = Linear(hidden * 2, hidden)
            self.fc2_bn_co = BatchNorm1d(hidden)
            self.fc2_co = Linear(hidden, hidden_out)
        elif self.args.cat_or_add == "add":
            self.fc1_bn_co = BatchNorm1d(hidden)
            self.fc1_co = Linear(hidden, hidden)
            self.fc2_bn_co = BatchNorm1d(hidden)
            self.fc2_co = Linear(hidden, hidden_out)
        else:
            assert False
        # 这部分代码用于对所有批量归一化层进行初始化。将权重初始化为常数1，偏置初始化为一个很小的非零常数（例如0.0001）
        # BN initialization.
        for m in self.modules():
            if isinstance(m, (torch.nn.BatchNorm1d)):
                torch.nn.init.constant_(m.weight, 1)
                torch.nn.init.constant_(m.bias, 0.0001)

    # 定义了模型的前向传播方法，接收输入数据 data 和一个布尔参数 eval_random，用于控制是否进行随机特征的评估。
    def forward(self, data, eval_random=True):
        # 如果输入数据中存在 x 属性，则将其赋给变量 x，否则将 data.feat 赋给 x。这里假设 x 包含节点特征。
        x = data.x if data.x is not None else data.feat
        # 从输入数据中获取边的索引和批次信息，其中 edge_index 存储了图的边的连接关系，而 batch 存储了每个节点所属的图批次。
        edge_index, batch = data.edge_index, data.batch
        # edge_index 是一个包含两行的张量，每一列代表一条边。第一行存储源节点的索引，第二行存储目标节点的索引。
        # row 变量存储了所有边的源节点索引，即第一行的内容。
        # col 变量存储了所有边的目标节点索引，即第二行的内容。
        row, col = edge_index
        # 将节点特征 x 经过批量归一化层处理，以规范化输入特征
        x = self.bn_feat(x)
        # 对输入特征 x 应用第一个图卷积层 self.conv_feat，使用 edge_index 定义的边连接关系。
        # F.relu(...)：对卷积层的输出应用 ReLU 激活函数。
        x = F.relu(self.conv_feat(x, edge_index))
        # 遍历图卷积层列表 self.convs 中的每一层。
        # x = self.bns_conv[i](x)：对当前图卷积层的输出进行批量归一化处理。
        # x = F.relu(conv(x, edge_index))：对经过批量归一化的输出应用 ReLU 激活函数
        for i, conv in enumerate(self.convs):
            x = self.bns_conv[i](x)
            x = F.relu(conv(x, edge_index))
        # 将每条边连接的两个节点的特征 x[row] 和 x[col] 沿着特征维度进行拼接，得到每条边的表示。shape(E,2D),E为边的数量，D是特征纬度
        edge_rep = torch.cat([x[row], x[col]], dim=-1)
        # 如果不使用边注意力机制，则将边权重初始化为均匀分布。shape(E,2)
        if self.without_edge_attention:
            edge_att = 0.5 * torch.ones(edge_rep.shape[0], 2).cuda()
        # 否则，将边的表示 edge_rep 经过线性变换并应用 softmax 函数计算边权重。shape(E,2)
        else:
            # 首先，edge_rep 作为输入被传递给 self.edge_att_mlp，这将对边的特征进行线性变换，并产生一个形状为(E,2) 的输出，其中每一行代表一个边，两个元素表示该边的两个注意力权重。
            # 然后，F.softmax 函数被应用于这个输出张量，指定 dim=-1 表示沿着最后一个维度进行 softmax 归一化。这将确保每一行的两个元素分别归一化为概率值，代表了该边在两种不同的注意力权重下的概率分布
            edge_att = F.softmax(self.edge_att_mlp(edge_rep), dim=-1)
        # 分别获取上下文边注意力和对象边注意力的边权重
        edge_weight_c = edge_att[:, 0]
        edge_weight_o = edge_att[:, 1]

        # 首先，根据 without_node_attention 参数的设置，决定是否计算节点的注意力权重。如果不需要节点注意力，则将节点注意力设置为均匀分布，即每个节点的注意力权重为 0.5。
        # 否则，通过神经网络 self.node_att_mlp 计算节点注意力权重，并进行 softmax 归一化。
        if self.without_node_attention:
            node_att = 0.5 * torch.ones(x.shape[0], 2).cuda()
        else:
            node_att = F.softmax(self.node_att_mlp(x), dim=-1)
        # 将节点特征 x 按照节点注意力权重进行加权求和，分别得到上下文节点的加权特征 xc 和对象节点的加权特征 xo。
        xc = node_att[:, 0].view(-1, 1) * x
        xo = node_att[:, 1].view(-1, 1) * x

        # 将上下文节点的加权特征和对象节点的加权特征分别输入上下文卷积层 self.context_convs 和对象卷积层 self.objects_convs 中进行信息传递，
        # 并使用 edge_weight_c 和 edge_weight_o 来对边进行加权
        xc = F.relu(self.context_convs(self.bnc(xc), edge_index, edge_weight_c))
        xo = F.relu(self.objects_convs(self.bno(xo), edge_index, edge_weight_o))

        # 对经过卷积层处理后的特征进行全局池化操作，将每个节点的特征汇总为图的特征，其中 self.global_pool 是全局池化函数
        xc = self.global_pool(xc, batch)
        xo = self.global_pool(xo, batch)

        # 将全局池化后的上下文节点特征 xc 和对象节点特征 xo 分别输入到上下文读出层 self.context_readout_layer 和对象读出层 self.objects_readout_layer 中，进行最终的分类预测
        xc_logis = self.context_readout_layer(xc)
        xo_logis = self.objects_readout_layer(xo)
        # 如果模型需要考虑随机特征，则调用 self.random_readout_layer 对上下文和对象节点的特征进行随机组合，并进行分类预测。
        xco_logis = self.random_readout_layer(xc, xo, eval_random=eval_random)

        # 最终返回上下文节点的分类预测结果 xc_logis、对象节点的分类预测结果 xo_logis，以及随机特征的分类预测结果 xco_logis。
        return xc_logis, xo_logis, xco_logis

    # 用于上下文节点的读出层，它对上下文节点进行进一步的特征处理和分类预测。
    def context_readout_layer(self, x):
        # x 通过批量归一化层 fc1_bn_c 进行批量归一化处理。
        x = self.fc1_bn_c(x)
        # 批量归一化后的特征 x 通过线性变换层 fc1_c 进行线性变换
        x = self.fc1_c(x)
        # 通过 ReLU 激活函数对线性变换后的特征进行激活，增加模型的非线性
        x = F.relu(x)
        # 对激活后的特征再次进行批量归一化处理
        x = self.fc2_bn_c(x)
        # 批量归一化后的特征再通过线性变换层 fc2_c 进行线性变换
        x = self.fc2_c(x)
        # 最后，对经过线性变换的特征进行 softmax 操作，得到类别预测的对数概率 x_logis。
        # 这里采用对数概率的形式输出，方便后续计算交叉熵损失。dim=-1 表示在最后一个维度（即特征维度）上进行 softmax 计算
        x_logis = F.log_softmax(x, dim=-1)
        return x_logis

    def objects_readout_layer(self, x):
   
        x = self.fc1_bn_o(x)
        x = self.fc1_o(x)
        x = F.relu(x)
        x = self.fc2_bn_o(x)
        x = self.fc2_o(x)
        x_logis = F.log_softmax(x, dim=-1)
        return x_logis

    # random_readout_layer() 方法是用于随机节点的读出层，它对随机节点进行进一步的特征处理和分类预测
    # 它接受三个参数：上下文节点特征 xc，对象节点特征 xo，以及一个布尔值 eval_random，用于控制是否进行随机操作
    def random_readout_layer(self, xc, xo, eval_random):
        # 获取上下文节点特征的数量，即节点的个数
        num = xc.shape[0]
        # 创建一个列表 l，其中包含从 0 到 num-1 的所有整数，用于后续随机操作
        l = [i for i in range(num)]
        # 如果模型允许随机操作，并且 eval_random 参数为 True，则对列表 l 进行随机打乱操作
        if self.with_random:
            if eval_random:
                random.shuffle(l)
        # 将打乱后的列表转换为 PyTorch 张量 random_idx，用于后续获取随机顺序的节点索引
        random_idx = torch.tensor(l)
        # 如果模型选择进行拼接操作，则将打乱后的上下文节点特征 xc 和对象节点特征 xo 沿着特征维度拼接在一起，得到新的特征表示 x
        if self.args.cat_or_add == "cat":
            x = torch.cat((xc[random_idx], xo), dim=1)
        else:
            # 如果模型选择进行相加操作，则将打乱后的上下文节点特征 xc 和对象节点特征 xo 相加，得到新的特征表示 x
            x = xc[random_idx] + xo

        x = self.fc1_bn_co(x)
        x = self.fc1_co(x)
        x = F.relu(x)
        x = self.fc2_bn_co(x)
        x = self.fc2_co(x)
        x_logis = F.log_softmax(x, dim=-1)
        return x_logis

class CausalGIN(torch.nn.Module):
    """GCN with BN and residual connection."""
    def __init__(self, num_features,
                       num_classes, args,
                gfn=False,
                edge_norm=True):
        super(CausalGIN, self).__init__()

        hidden = args.hidden
        num_conv_layers = args.layers
        self.args = args
        self.global_pool = global_add_pool
        GConv = partial(GCNConv, edge_norm=edge_norm, gfn=gfn)
        hidden_in = num_features
        self.num_classes = num_classes
        hidden_out = num_classes
        self.fc_num = args.fc_num
        self.bn_feat = BatchNorm1d(hidden_in)
        self.conv_feat = GCNConv(hidden_in, hidden, gfn=True) # linear transform
        self.bns_conv = torch.nn.ModuleList()
        self.convs = torch.nn.ModuleList()
        for i in range(num_conv_layers):
            self.convs.append(GINConv(
            Sequential(
                       Linear(hidden, hidden), 
                       BatchNorm1d(hidden), 
                       ReLU(),
                       Linear(hidden, hidden), 
                       ReLU())))

        self.edge_att_mlp = nn.Linear(hidden * 2, 2)
        self.node_att_mlp = nn.Linear(hidden, 2)
        self.bnc = BatchNorm1d(hidden)
        self.bno= BatchNorm1d(hidden)
        self.context_convs = GConv(hidden, hidden)
        self.objects_convs = GConv(hidden, hidden)

        # context mlp
        self.fc1_bn_c = BatchNorm1d(hidden)
        self.fc1_c = Linear(hidden, hidden)
        self.fc2_bn_c = BatchNorm1d(hidden)
        self.fc2_c = Linear(hidden, hidden_out)
        # object mlp
        self.fc1_bn_o = BatchNorm1d(hidden)
        self.fc1_o = Linear(hidden, hidden)
        self.fc2_bn_o = BatchNorm1d(hidden)
        self.fc2_o = Linear(hidden, hidden_out)
        # random mlp
        if self.args.cat_or_add == "cat":
            self.fc1_bn_co = BatchNorm1d(hidden * 2)
            self.fc1_co = Linear(hidden * 2, hidden)
            self.fc2_bn_co = BatchNorm1d(hidden)
            self.fc2_co = Linear(hidden, hidden_out)

        elif self.args.cat_or_add == "add":
            self.fc1_bn_co = BatchNorm1d(hidden)
            self.fc1_co = Linear(hidden, hidden)
            self.fc2_bn_co = BatchNorm1d(hidden)
            self.fc2_co = Linear(hidden, hidden_out)
        else:
            assert False
        
        # BN initialization.
        for m in self.modules():
            if isinstance(m, (torch.nn.BatchNorm1d)):
                torch.nn.init.constant_(m.weight, 1)
                torch.nn.init.constant_(m.bias, 0.0001)

    def forward(self, data, eval_random=True, train_type="base"):

        x = data.x if data.x is not None else data.feat
        edge_index, batch = data.edge_index, data.batch
        row, col = edge_index
        x = self.bn_feat(x)
        x = F.relu(self.conv_feat(x, edge_index))
        
        for i, conv in enumerate(self.convs):
            x = conv(x, edge_index)
        
        edge_rep = torch.cat([x[row], x[col]], dim=-1)
        edge_att = F.softmax(self.edge_att_mlp(edge_rep), dim=-1)
        edge_weight_c = edge_att[:, 0]
        edge_weight_o = edge_att[:, 1]

        node_att = F.softmax(self.node_att_mlp(x), dim=-1)
        node_weight_c = node_att[:, 0]
        node_weight_o = node_att[:, 1]
        
        
        xc = node_weight_c.view(-1, 1) * x
        xo = node_weight_o.view(-1, 1) * x
        xc = F.relu(self.context_convs(self.bnc(xc), edge_index, edge_weight_c))
        xo = F.relu(self.objects_convs(self.bno(xo), edge_index, edge_weight_o))

        xc = self.global_pool(xc, batch)
        xo = self.global_pool(xo, batch)
        
        xc_logis = self.context_readout_layer(xc)
        xco_logis = self.random_readout_layer(xc, xo, eval_random=eval_random)
        # return xc_logis, xo_logis, xco_logis
        xo_logis = self.objects_readout_layer(xo, train_type)
        return xc_logis, xo_logis, xco_logis
        


    def context_readout_layer(self, x):
        
        x = self.fc1_bn_c(x)
        x = self.fc1_c(x)
        x = F.relu(x)
        x = self.fc2_bn_c(x)
        x = self.fc2_c(x)
        x_logis = F.log_softmax(x, dim=-1)
        return x_logis

    def objects_readout_layer(self, x, train_type):
   
        x = self.fc1_bn_o(x)
        x = self.fc1_o(x)
        x = F.relu(x)
        x = self.fc2_bn_o(x)
        x = self.fc2_o(x)
        x_logis = F.log_softmax(x, dim=-1)
        if train_type == "irm":
            return x, x_logis
        else:
            return x_logis

    def random_readout_layer(self, xc, xo, eval_random):

        num = xc.shape[0]
        l = [i for i in range(num)]
        if eval_random:
            random.shuffle(l)
        random_idx = torch.tensor(l)
        
        if self.args.cat_or_add == "cat":
            x = torch.cat((xc[random_idx], xo), dim=1)
        else:
            x = xc[random_idx] + xo

        x = self.fc1_bn_co(x)
        x = self.fc1_co(x)
        x = F.relu(x)
        x = self.fc2_bn_co(x)
        x = self.fc2_co(x)
        x_logis = F.log_softmax(x, dim=-1)
        return x_logis

class CausalGAT(torch.nn.Module):
    def __init__(self, num_features,
                       num_classes, 
                       args, 
                       head=4, 
                       dropout=0.2):
        super(CausalGAT, self).__init__()
        num_conv_layers = args.layers
        hidden = args.hidden
        self.args = args
        self.global_pool = global_add_pool
        self.dropout = dropout
        GConv = partial(GCNConv, edge_norm=True, gfn=False)

        hidden_in = num_features
        self.num_classes = num_classes
        hidden_out = num_classes
        self.fc_num = args.fc_num
        self.bn_feat = BatchNorm1d(hidden_in)
        self.conv_feat = GCNConv(hidden_in, hidden, gfn=True) # linear transform
        self.bns_conv = torch.nn.ModuleList()
        self.convs = torch.nn.ModuleList()

        for i in range(num_conv_layers):
            self.bns_conv.append(BatchNorm1d(hidden))
            self.convs.append(GATConv(hidden, int(hidden / head), heads=head, dropout=dropout))

        self.edge_att_mlp = nn.Linear(hidden * 2, 2)
        self.node_att_mlp = nn.Linear(hidden, 2)
        self.bnc = BatchNorm1d(hidden)
        self.bno= BatchNorm1d(hidden)
        self.context_convs = GConv(hidden, hidden)
        self.objects_convs = GConv(hidden, hidden)

        # context mlp
        self.fc1_bn_c = BatchNorm1d(hidden)
        self.fc1_c = Linear(hidden, hidden)
        self.fc2_bn_c = BatchNorm1d(hidden)
        self.fc2_c = Linear(hidden, hidden_out)
        # object mlp
        self.fc1_bn_o = BatchNorm1d(hidden)
        self.fc1_o = Linear(hidden, hidden)
        self.fc2_bn_o = BatchNorm1d(hidden)
        self.fc2_o = Linear(hidden, hidden_out)
        # random mlp
        if self.args.cat_or_add == "cat":
            self.fc1_bn_co = BatchNorm1d(hidden * 2)
            self.fc1_co = Linear(hidden * 2, hidden)
            self.fc2_bn_co = BatchNorm1d(hidden)
            self.fc2_co = Linear(hidden, hidden_out)

        elif self.args.cat_or_add == "add":
            self.fc1_bn_co = BatchNorm1d(hidden)
            self.fc1_co = Linear(hidden, hidden)
            self.fc2_bn_co = BatchNorm1d(hidden)
            self.fc2_co = Linear(hidden, hidden_out)
        else:
            assert False
        
        # BN initialization.
        for m in self.modules():
            if isinstance(m, (torch.nn.BatchNorm1d)):
                torch.nn.init.constant_(m.weight, 1)
                torch.nn.init.constant_(m.bias, 0.0001)

    def forward(self, data, eval_random=True):

        x = data.x if data.x is not None else data.feat
        edge_index, batch = data.edge_index, data.batch
        row, col = edge_index
        x = self.bn_feat(x)
        x = F.relu(self.conv_feat(x, edge_index))
        
        for i, conv in enumerate(self.convs):
            x = self.bns_conv[i](x)
            x = F.relu(conv(x, edge_index))
        
        edge_rep = torch.cat([x[row], x[col]], dim=-1)
        edge_att = F.softmax(self.edge_att_mlp(edge_rep), dim=-1)
        edge_weight_c = edge_att[:, 0]
        edge_weight_o = edge_att[:, 1]

        node_att = F.softmax(self.node_att_mlp(x), dim=-1)
        xc = node_att[:, 0].view(-1, 1) * x
        xo = node_att[:, 1].view(-1, 1) * x
        xc = F.relu(self.context_convs(self.bnc(xc), edge_index, edge_weight_c))
        xo = F.relu(self.objects_convs(self.bno(xo), edge_index, edge_weight_o))

        xc = self.global_pool(xc, batch)
        xo = self.global_pool(xo, batch)
        
        xc_logis = self.context_readout_layer(xc)
        xo_logis = self.objects_readout_layer(xo)
        xco_logis = self.random_readout_layer(xc, xo, eval_random=eval_random)
        return xc_logis, xo_logis, xco_logis

    def context_readout_layer(self, x):
        
        x = self.fc1_bn_c(x)
        x = self.fc1_c(x)
        x = F.relu(x)
        x = self.fc2_bn_c(x)
        x = self.fc2_c(x)
        x_logis = F.log_softmax(x, dim=-1)
        return x_logis

    def objects_readout_layer(self, x):
   
        x = self.fc1_bn_o(x)
        x = self.fc1_o(x)
        x = F.relu(x)
        x = self.fc2_bn_o(x)
        x = self.fc2_o(x)
        x_logis = F.log_softmax(x, dim=-1)
        return x_logis

    def random_readout_layer(self, xc, xo, eval_random):

        num = xc.shape[0]
        l = [i for i in range(num)]
        if eval_random:
            random.shuffle(l)
        random_idx = torch.tensor(l)
        
        if self.args.cat_or_add == "cat":
            x = torch.cat((xc[random_idx], xo), dim=1)
        else:
            x = xc[random_idx] + xo

        x = self.fc1_bn_co(x)
        x = self.fc1_co(x)
        x = F.relu(x)
        x = self.fc2_bn_co(x)
        x = self.fc2_co(x)
        x_logis = F.log_softmax(x, dim=-1)
        return x_logis

class GCNNet(torch.nn.Module):
    """GCN with BN and residual connection."""
    def __init__(self, num_features,
                       num_classes, hidden, 
                       num_feat_layers=1, 
                       num_conv_layers=3,
                 num_fc_layers=2, gfn=False, collapse=False, residual=False,
                 res_branch="BNConvReLU", global_pool="sum", dropout=0, 
                 edge_norm=True):
        super(GCNNet, self).__init__()

        # 初始化全局池化函数和dropout参数
        self.global_pool = global_add_pool
        self.dropout = dropout
        # # 使用偏函数定义GConv，以便后续创建GCNConv层
        GConv = partial(GCNConv, edge_norm=edge_norm, gfn=gfn)
        # # 定义输入特征的维度
        hidden_in = num_features

        # 初始化特征层和卷积层
        self.bn_feat = BatchNorm1d(hidden_in)
        self.conv_feat = GCNConv(hidden_in, hidden, gfn=True) # linear transform
        self.bns_conv = torch.nn.ModuleList()
        self.convs = torch.nn.ModuleList()

        # 初始化多层图卷积层
        for i in range(num_conv_layers):
            self.bns_conv.append(BatchNorm1d(hidden))
            self.convs.append(GConv(hidden, hidden))

        # 初始化隐藏层的批量归一化层和线性层
        self.bn_hidden = BatchNorm1d(hidden)
        self.bns_fc = torch.nn.ModuleList()
        self.lins = torch.nn.ModuleList()

        # 初始化多层全连接层
        for i in range(num_fc_layers - 1):
            self.bns_fc.append(BatchNorm1d(hidden))
            self.lins.append(Linear(hidden, hidden))

        # 初始化输出层的线性层
        self.lin_class = Linear(hidden, num_classes)

        # BN initialization.
        for m in self.modules():
            if isinstance(m, (torch.nn.BatchNorm1d)):
                torch.nn.init.constant_(m.weight, 1)
                torch.nn.init.constant_(m.bias, 0.0001)

    def forward(self, data):
        # 获取输入数据的特征和边索引
        x = data.x if data.x is not None else data.feat
        edge_index, batch = data.edge_index, data.batch

        # 对输入特征进行批量归一化和ReLU激活
        x = self.bn_feat(x)
        x = F.relu(self.conv_feat(x, edge_index))

        # 逐层进行图卷积操作，同时对结果进行批量归一化和ReLU激活
        for i, conv in enumerate(self.convs):
            x = self.bns_conv[i](x)
            x = F.relu(conv(x, edge_index))

        # 对图节点特征进行全局池化
        x = self.global_pool(x, batch)

        # 逐层进行全连接操作，同时对结果进行批量归一化和ReLU激活
        for i, lin in enumerate(self.lins):
            x = self.bns_fc[i](x)
            x = F.relu(lin(x))

        # 对隐藏层结果进行批量归一化
        x = self.bn_hidden(x)

        # 输出层线性变换，并进行log_softmax激活
        x = self.lin_class(x)
        return F.log_softmax(x, dim=-1)

class GINNet(torch.nn.Module):
    def __init__(self, num_features,
                       num_classes,
                       hidden, 
                       num_fc_layers=2, 
                       num_conv_layers=3, 
                       dropout=0):

        super(GINNet, self).__init__()
        self.global_pool = global_add_pool
        self.dropout = dropout
        hidden_in = num_features
        hidden_out = num_classes
        
        self.bn_feat = BatchNorm1d(hidden_in)
        self.conv_feat = GCNConv(hidden_in, hidden, gfn=True) # linear transform
        
        self.convs = torch.nn.ModuleList()
        for i in range(num_conv_layers):
            self.convs.append(GINConv(
            Sequential(Linear(hidden, hidden), 
                       BatchNorm1d(hidden), 
                       ReLU(),
                       Linear(hidden, hidden), 
                       ReLU())))

        self.bn_hidden = BatchNorm1d(hidden)
        self.bns_fc = torch.nn.ModuleList()
        self.lins = torch.nn.ModuleList()

        for i in range(num_fc_layers - 1):
            self.bns_fc.append(BatchNorm1d(hidden))
            self.lins.append(Linear(hidden, hidden))
        self.lin_class = Linear(hidden, hidden_out)

        # BN initialization.
        for m in self.modules():
            if isinstance(m, (torch.nn.BatchNorm1d)):
                torch.nn.init.constant_(m.weight, 1)
                torch.nn.init.constant_(m.bias, 0.0001)

    def forward(self, data):
        x = data.x if data.x is not None else data.feat
        edge_index, batch = data.edge_index, data.batch
        # x, edge_index, batch = data.feat, data.edge_index, data.batch
        x = self.bn_feat(x)
        x = F.relu(self.conv_feat(x, edge_index))
        
        for i, conv in enumerate(self.convs):
            x = conv(x, edge_index)
            
        x = self.global_pool(x, batch)
        for i, lin in enumerate(self.lins):
            x = self.bns_fc[i](x)
            x = F.relu(lin(x))    
        x = self.bn_hidden(x)
        x = self.lin_class(x)

        prediction = F.log_softmax(x, dim=-1)
        return prediction

class GATNet(torch.nn.Module):
    def __init__(self, num_features, 
                       num_classes,
                       hidden,
                       head=4,
                       num_fc_layers=2, 
                       num_conv_layers=3, 
                       dropout=0.2):

        super(GATNet, self).__init__()

        self.global_pool = global_add_pool
        self.dropout = dropout
        hidden_in = num_features
        hidden_out = num_classes
   
        self.bn_feat = BatchNorm1d(hidden_in)
        self.conv_feat = GCNConv(hidden_in, hidden, gfn=True) # linear transform
        self.bns_conv = torch.nn.ModuleList()
        self.convs = torch.nn.ModuleList()

        for i in range(num_conv_layers):
            self.bns_conv.append(BatchNorm1d(hidden))
            self.convs.append(GATConv(hidden, int(hidden / head), heads=head, dropout=dropout))
        self.bn_hidden = BatchNorm1d(hidden)
        self.bns_fc = torch.nn.ModuleList()
        self.lins = torch.nn.ModuleList()

        for i in range(num_fc_layers - 1):
            self.bns_fc.append(BatchNorm1d(hidden))
            self.lins.append(Linear(hidden, hidden))
        self.lin_class = Linear(hidden, hidden_out)

        # BN initialization.
        for m in self.modules():
            if isinstance(m, (torch.nn.BatchNorm1d)):
                torch.nn.init.constant_(m.weight, 1)
                torch.nn.init.constant_(m.bias, 0.0001)

    def forward(self, data):
        
        x = data.x if data.x is not None else data.feat
        edge_index, batch = data.edge_index, data.batch
        
        x = self.bn_feat(x)
        x = F.relu(self.conv_feat(x, edge_index))
        
        for i, conv in enumerate(self.convs):
            x = self.bns_conv[i](x)
            x = F.relu(conv(x, edge_index))

        x = self.global_pool(x, batch)
        for i, lin in enumerate(self.lins):
            x = self.bns_fc[i](x)
            x = F.relu(lin(x))

        x = self.bn_hidden(x)
        if self.dropout > 0:
            x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.lin_class(x)
        return F.log_softmax(x, dim=-1)
